import base64
import pickle
import threading
from box_util import box_iou_xyxy
import time
import data_class

import numpy as np
from YOLOv3 import YOLOv3
import paddle
from img_loader import single_image_data_loader
from nms import nms
from draw_rec import draw_results


ANCHORS = [10, 13, 16, 30, 33, 23, 30, 61, 62,
           45, 59, 119, 116, 90, 156, 198, 373, 326]
ANCHOR_MASKS = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]
VALID_THRESH = 0.01
NMS_TOPK = 400
NMS_POSK = 100
NMS_THRESH = 0.45

NUM_CLASSES = 7
try:
    MODEL_STATE_DICT = paddle.load(r'../params/yolo_epoch50.pdparams')  # 模型参数
    model = YOLOv3(num_classes=NUM_CLASSES)
    model.load_dict(MODEL_STATE_DICT)
    model.eval()
except(ValueError):
    pass


# 加载图片，这一步只能在客户端做
def load_img(img_path: str):
    test_loader = single_image_data_loader(img_path, mode='test')
    data_generator = test_loader()  # test_loader()返回一个生成器
    img_name, img_data, img_scale_data = next(
        data_generator)  # img_data就是一个多维ndarray
    return img_name, img_data, img_scale_data


class MainService:
    __instance = None
    # 初始化锁
    __init_lock = threading.Lock()

    def __init__(self) -> None:
        pass

    @classmethod
    def get_instance(cls):
        with MainService.__init_lock:
            if not cls.__instance:
                cls.__instance = MainService()
            return cls.__instance

    # 特征提取
    def feature_extraction(self, img_data):

        img = paddle.to_tensor(np.array(img_data))
        outputs = model.forward(img)  # outputs是一个ndarray的list
        return outputs

    # 获取预测框以及得分 参数outputs是一个ndarray的list. 函数返回的都是ndarray的list
    def get_bboxes_and_scores(self, im_shape, prev_step_output):

        im_shape = paddle.to_tensor(np.array(im_shape))
        downsample = 32
        total_boxes = []
        total_scores = []
        for i, out in enumerate(prev_step_output):
            out = paddle.to_tensor(out)
            anchor_mask = ANCHOR_MASKS[i]
            anchors_this_level = []
            for m in anchor_mask:
                anchors_this_level.append(ANCHORS[2 * m])
                anchors_this_level.append(ANCHORS[2 * m + 1])

            boxes, scores = paddle.vision.ops.yolo_box(
                x=out,
                img_size=im_shape,
                anchors=anchors_this_level,
                class_num=NUM_CLASSES,
                conf_thresh=VALID_THRESH,
                downsample_ratio=downsample,
                name="yolo_box" + str(i))

            scores = paddle.transpose(scores, perm=[0, 2, 1])
            total_boxes.append(boxes.numpy())
            total_scores.append(scores.numpy())
            downsample = downsample // 2

        return total_boxes, total_scores

    # 拼接不同层级的结果 total_boxes和total_scores都是ndarray的list. 函数的返回内容也都是ndarray的list
    def concat_result(self, total_boxes, total_scores):

        for i in range(len(total_boxes)):
            total_boxes[i] = paddle.to_tensor(total_boxes[i])
        # for i in range(len(con_resultitem.total_scores)):
        for i in range(len(total_scores)):
            total_scores[i] = paddle.to_tensor(total_scores[i])
        yolo_boxes = paddle.concat(total_boxes, axis=1)
        yolo_scores = paddle.concat(total_scores, axis=2)
        yolo_boxes = yolo_boxes.numpy()
        yolo_scores = yolo_scores.numpy()

        return yolo_boxes, yolo_scores

    def multiclass_nms(self, bboxes, scores):
        """
        This is for multiclass_nms
        """

        score_thresh = 0.01
        nms_thresh = 0.45
        pre_nms_topk = 1000
        pos_nms_topk = 100
        batch_size = bboxes.shape[0]
        class_num = scores.shape[1]
        rets = []
        for i in range(batch_size):
            bboxes_i = bboxes[i]
            scores_i = scores[i]
            ret = []
            for c in range(class_num):
                scores_i_c = scores_i[c]
                keep_inds = nms(bboxes_i, scores_i_c, score_thresh,
                                nms_thresh, pre_nms_topk, i=i, c=c)
                if len(keep_inds) < 1:
                    continue
                keep_bboxes = bboxes_i[keep_inds]
                keep_scores = scores_i_c[keep_inds]
                keep_results = np.zeros([keep_scores.shape[0], 6])
                keep_results[:, 0] = c
                keep_results[:, 1] = keep_scores[:]
                keep_results[:, 2:6] = keep_bboxes[:, :]
                ret.append(keep_results)
            if len(ret) < 1:
                rets.append(ret)
                continue
            ret_i = np.concatenate(ret, axis=0)
            scores_i = ret_i[:, 1]
            if len(scores_i) > pos_nms_topk:
                inds = np.argsort(scores_i)[::-1]
                inds = inds[:pos_nms_topk]
                ret_i = ret_i[inds]

            rets.append(ret_i)

        return rets


# 每个步骤的用时
load_img_times = []
feature_extraction_times = []
bboxes_and_scores_times = []
concat_times = []
nms_tims = []
show_res_times = []


# 以下方法是不能寻址的，用作测试

# 加载图片，这一步只能在客户端做
def load_img(img_path: str):
    test_loader = single_image_data_loader(img_path, mode='test')
    data_generator = test_loader()  # test_loader()返回一个生成器
    img_name, img_data, img_scale_data = next(
        data_generator)  # img_data就是一个多维ndarray
    return img_name, img_data, img_scale_data


# 特征提取
def feature_extraction(img_data: np.ndarray):
    img = paddle.to_tensor(img_data)
    outputs = model.forward(img)  # outputs是一个ndarray的list
    return outputs


# 获取预测框以及得分 参数outputs是一个ndarray的list. 函数返回的都是ndarray的list
def get_bboxes_and_scores(outputs: list, im_shape: np.ndarray):
    im_shape = paddle.to_tensor(im_shape)
    downsample = 32
    total_boxes = []
    total_scores = []
    for i, out in enumerate(outputs):
        out = paddle.to_tensor(out)
        anchor_mask = ANCHOR_MASKS[i]
        anchors_this_level = []
        for m in anchor_mask:
            anchors_this_level.append(ANCHORS[2 * m])
            anchors_this_level.append(ANCHORS[2 * m + 1])

        boxes, scores = paddle.vision.ops.yolo_box(
            x=out,
            img_size=im_shape,
            anchors=anchors_this_level,
            class_num=NUM_CLASSES,
            conf_thresh=VALID_THRESH,
            downsample_ratio=downsample,
            name="yolo_box" + str(i))

        scores = paddle.transpose(scores, perm=[0, 2, 1])
        total_boxes.append(boxes.numpy())
        total_scores.append(scores.numpy())
        downsample = downsample // 2

    return total_boxes, total_scores


# 拼接不同层级的结果 total_boxes和total_scores都是ndarray的list. 函数的返回内容也都是ndarray的list
def concat_result(total_boxes: list, total_scores: list):
    for i in range(len(total_boxes)):
        total_boxes[i] = paddle.to_tensor(total_boxes[i])
    for i in range(len(total_scores)):
        total_scores[i] = paddle.to_tensor(total_scores[i])
    yolo_boxes = paddle.concat(total_boxes, axis=1)
    yolo_scores = paddle.concat(total_scores, axis=2)
    yolo_boxes = yolo_boxes.numpy()
    yolo_scores = yolo_scores.numpy()
    return yolo_boxes, yolo_scores


# 多分类非极大值抑制
def multiclass_nms(bboxes, scores, score_thresh=0.01, nms_thresh=0.45, pre_nms_topk=1000, pos_nms_topk=100):
    """
    This is for multiclass_nms
    """
    batch_size = bboxes.shape[0]
    class_num = scores.shape[1]
    rets = []
    for i in range(batch_size):
        bboxes_i = bboxes[i]
        scores_i = scores[i]
        ret = []
        for c in range(class_num):
            scores_i_c = scores_i[c]
            keep_inds = nms(bboxes_i, scores_i_c, score_thresh,
                            nms_thresh, pre_nms_topk, i=i, c=c)
            if len(keep_inds) < 1:
                continue
            keep_bboxes = bboxes_i[keep_inds]
            keep_scores = scores_i_c[keep_inds]
            keep_results = np.zeros([keep_scores.shape[0], 6])
            keep_results[:, 0] = c
            keep_results[:, 1] = keep_scores[:]
            keep_results[:, 2:6] = keep_bboxes[:, :]
            ret.append(keep_results)
        if len(ret) < 1:
            rets.append(ret)
            continue
        ret_i = np.concatenate(ret, axis=0)
        scores_i = ret_i[:, 1]
        if len(scores_i) > pos_nms_topk:
            inds = np.argsort(scores_i)[::-1]
            inds = inds[:pos_nms_topk]
            ret_i = ret_i[inds]

        rets.append(ret_i)

    return rets[0]


def predict(img_path: str):
    # 加载图片
    start_time = time.time()
    _, img_data, img_scale = load_img(img_path)
    load_img_times.append(time.time() - start_time)

    # 特征提取
    # outputs是一个ndarray的list
    start_time = time.time()
    outputs = feature_extraction(img_data)
    feature_extraction_times.append(time.time() - start_time)

    # 预测框以及得分
    start_time = time.time()
    bboxes, scores = get_bboxes_and_scores(outputs, im_shape=img_scale)
    bboxes_and_scores_times.append(time.time() - start_time)

    # 拼接预测框以及得分
    start_time = time.time()
    concated_box, concated_score = concat_result(bboxes, scores)
    concat_times.append(time.time() - start_time)

    # 多分类非极大值抑制(NMS),在multi_nms.py中定义. result是一个ndarray
    start_time = time.time()
    result = multiclass_nms(concated_box, concated_score,
                            score_thresh=VALID_THRESH,
                            nms_thresh=NMS_THRESH,
                            pre_nms_topk=NMS_TOPK,
                            pos_nms_topk=NMS_POSK)
    nms_tims.append(time.time() - start_time)

    # 显示结果,见draw_rec.py. 这一步必需要在客户端
    # 这里为了计时，把显示图片取消了。如果要显示图片的话，在draw_rec.py中把最后一行解注释即可
    start_time = time.time()
    draw_results(result, img_path, draw_thresh=0.5)
    show_res_times.append(time.time() - start_time)


# 输出花费的时间
def print_spend_times():
    load_avg = np.mean(load_img_times)
    feature_avg = np.mean(feature_extraction_times)
    box_score_avg = np.mean(bboxes_and_scores_times)
    concat_avg = np.mean(concat_times)
    nms_avg = np.mean(nms_tims)
    show_res_avg = np.mean(show_res_times)
    print("加载图片平均用时: {}".format(load_avg))
    print("特征提取平均用时: {}".format(feature_avg))
    print("获取预测框以及得分平均用时: {}".format(box_score_avg))
    print("拼接预测框以及得分平均用时: {}".format(concat_avg))
    print("NMS平均用时: {}".format(nms_avg))
    print("输出结果平均用时: {}".format(show_res_avg))
    print("总平均用时: {}".format(load_avg + feature_avg +
          box_score_avg + concat_avg + nms_avg + show_res_avg))

# if __name__ == '__main__':
#     image_name = r'../../client/test_img/1851.jpeg'
#     test_times = 10  # 测试10次
#     for i in range(0, test_times):
#         print('{}/{}'.format(i + 1, test_times))
#         predict(image_name)
#     print_spend_times()


main_service = MainService.get_instance()
